// Copyright 2012 The golab Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package exp

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"sync"
	"time"

	"../cmd"
	"../machine"
	"../minjs"
)

type Exp struct {
	Log      *log.Logger         `json:"-"` // used to log on a per-experiment level
	ID       string              // a descriptive name for the experiment
	Machines *machine.Group      // the groups of machines to work on
	Cmds     map[string]*cmd.Cmd // the pool of commands to draw from
	Steps    []*Step             // the steps of the experiment to follow
	st       time.Time           // record the start time of the experiment
	as       map[string]*Step    // active steps that will need to be stopped
	mux      sync.RWMutex        // protect the as map to keep it thread safe
	rp       string              // results file path
	lf       *os.File            // output file for Exp.Log
}

// NewExp opens and parse an experiment from a file named by the string input.
func NewExp(path string) (*Exp, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	b, err := ioutil.ReadAll(f)
	if err != nil {
		return nil, err
	}
	return NewExpStr(b)
}

// NewExpStr parses an experiment from a []byte.
func NewExpStr(exp []byte) (*Exp, error) {
	e := &Exp{}
	err := json.Unmarshal(minjs.RmComments(exp), e)
	if err != nil {
		return nil, err
	}
	err = e.setup()
	if err != nil && e.lf != nil {
		if e.Log != nil {
			e.Log.Print("SETUP CRASH")
		}
		e.lf.Close()
	}
	return e, err
}

// setup does a lot of internal heavy lifting.
func (e *Exp) setup() error {
	e.st = time.Now()

	// Set the results path and make the directory.
	e.rp = fmt.Sprintf("%s/%d", e.ID, e.st.Unix())
	err := os.MkdirAll(e.rp, os.ModeDir|os.ModePerm)
	if err != nil {
		return err
	}

	// Open a global log file.
	e.lf, err = os.Create(fmt.Sprintf("%s/exp.log", e.rp))
	if err != nil {
		e, ok := err.(*os.PathError)
		if ok {
			return e.Err
		}
		return err
	}
	e.Log = log.New(e.lf, "", log.LstdFlags)
	e.Log.Printf("ok: setup: %s started", e.rp)

	// Save a copy of what we are running.
	efp := fmt.Sprintf("%s/exp.json", e.rp)
	e.Log.Printf("ok: setup: config file saved to %s", efp)

	e.as = make(map[string]*Step)
	if e.Machines.ID == "" {
		e.Machines.ID = "Machines"
	}
	e.Machines.FixIDs(e.Log, e.rp)
	for k, v := range e.Cmds {
		if v != nil {
			v.ID = k
		}
	}

	ef, err := os.Create(efp)
	if err != nil {
		e, ok := err.(*os.PathError)
		if ok {
			return e.Err
		}
		return err
	}
	b, err := json.MarshalIndent(e, "", "  ")
	if err != nil {
		ef.Close()
		return err
	}
	_, err = ef.Write(b)
	if err != nil {
		ef.Close()
		return err
	}
	err = ef.Sync()
	if err != nil {
		ef.Close()
		return err
	}
	return ef.Close()
}

// Run will take a parsed experiment and apply steps to groups of machines.
func (e *Exp) Run() error {
	var s *Step
	for i := range e.Steps {
		s = e.Steps[i]
		if s.ID == "" {
			// Assign a generated ID for the step if we don't have one.
			s.ID = fmt.Sprintf("%d", i)
		}
		// Catch reserved cmdIDs to stop running steps.
		if s.CmdID == "WAIT" {
			e.mux.RLock()
			p := e.as[s.ParentID]
			e.mux.RUnlock()
			if p == nil {
				e.Log.Printf("WARNING: Run: WAIT on nil skipped (ParentID=%s)", s.ParentID)
				continue
			}
			err := p.Wait(e)
			if err != nil {
				e.Log.Printf("ERROR: Run:p.Wait: %v (StepID=%s ParentID=%s)", err, s.ID, p.ID)
				e.crash()
				return err
			}
		} else if s.CmdID == "STOP" || s.CmdID == "KILL" {
			e.mux.RLock()
			p := e.as[s.ParentID]
			e.mux.RUnlock()
			if p == nil {
				e.Log.Printf("WARNING: Run: %s on nil skipped (ParentID=%s)", s.CmdID, s.ParentID)
				continue
			}
			if s.CmdID == "KILL" {
				p.Kill = true
			}
			err := p.Stop(e)
			if err != nil {
				e.Log.Printf("ERROR: Run:p.Stop: %v (StepID=%s ParentID=%s)", err, s.ID, p.ID)
				e.crash()
				return err
			}
		} else {
			// If we are not stopping then we should be starting a step.
			err := s.Start(e)
			if err != nil {
				e.Log.Printf("ERROR: Run:s.Start: %v (StepID=%s)", err, s.ID)
				e.crash()
				return err
			}
		}
		e.Log.Printf("ok: Run: (StepID=%s CmdID=%s)", s.ID, s.CmdID)
	}
	// Everything at this point has started, and all that is left is to clean up
	// remaining steps that are running.
	for i := len(e.Steps) - 1; i > -1; i-- {
		e.mux.RLock()
		s = e.as[e.Steps[i].ID]
		e.mux.RUnlock()
		if s == nil {
			continue
		}
		if s.Kill {
			err := s.Stop(e)
			if err != nil {
				e.Log.Printf("ERROR: Run:s.Stop: %v (StepID=%s CmdID=%s)", err, s.ID, s.CmdID)
				e.crash()
				return err
			}
			e.Log.Printf("ok: Run:s.Stop: (StepID=%s CmdID=%s)", s.ID, s.CmdID)
		} else {
			err := s.Wait(e)
			if err != nil {
				e.Log.Printf("ERROR: Run:s.Wait: %v (StepID=%s CmdID=%s)", err, s.ID, s.CmdID)
				e.crash()
				return err
			}
			e.Log.Printf("ok: Run:s.Wait: (StepID=%s CmdID=%s)", s.ID, s.CmdID)
		}
	}
	return nil
}

// Save loops through all the steps and saves the results if required.
func (e *Exp) Save() error {
	var s *Step
	for i := range e.Steps {
		s = e.Steps[i]
		if s.SaveRes {
			err := s.Save()
			if err != nil {
				e.Log.Printf("ERROR: Save: %v (StepID=%s CmdID=%s)", err, s.ID, s.CmdID)
				e.crash()
				return err
			}
			e.Log.Printf("ok: Save: (StepID=%s CmdID=%s)", s.ID, s.CmdID)
		}
	}
	return nil
}

// Close cleans up state and frees up resources for the experiment.
func (e *Exp) Close() error {
	e.Machines.Close(e.Log)
	e.Log.Print("CLOSING LOG FILE")
	return e.lf.Close()
}

// crash is used to try to gracefully close out an experiment when something
// goes wrong.
func (e *Exp) crash() {
	for i := range e.as {
		e.as[i].crash()
	}
	e.Log.Print("CRASH")
	e.lf.Close()
}

type Step struct {
	GroupID  string         // The group that will run the command.
	CmdID    string         // The command to run from the experiment's list of commands.
	TargetID string         `json:",omitempty"` // The group that will be the target of the command.
	Async    bool           `json:",omitempty"` // Should this command run asynchronously?
	Kill     bool           `json:",omitempty"` // Should we kill it instead of waiting for it to finish?
	ID       string         `json:",omitempty"` // Used to identify a step to be stopped later.
	ParentID string         `json:",omitempty"` // The step to wait/kill.
	SaveRes  bool           `json:",omitempty"` // Should we save the result to a file?
	g        *machine.Group // group that will run the step
	tg       *machine.Group // group that will be the target of the step
}

// Start creates and runs a subprocess on a group of machines.
func (s *Step) Start(e *Exp) error {
	s.g = e.Machines.Find(s.GroupID)
	if s.g == nil {
		return fmt.Errorf("GroupID=%s not found.", s.GroupID)
	}
	if s.TargetID != "" {
		//TODO(gavaletz) can we have a target that is not running RPC server.
		s.tg = e.Machines.Find(s.TargetID)
		if s.tg == nil {
			return fmt.Errorf("TargetID=%s not found.", s.GroupID)
		}
	}
	c := e.Cmds[s.CmdID]
	if c == nil {
		return fmt.Errorf("CmdID=%s not found.", s.CmdID)
	}
	var err error
	if s.tg == nil {
		err = s.g.Start(c)
	} else {
		err = s.g.StartTarget(c, s.tg)
	}
	if err != nil {
		return err
	}
	e.mux.Lock()
	e.as[s.ID] = s
	e.mux.Unlock()
	if s.Async {
		return err
	}
	return s.Wait(e)
}

// Wait will block on a subprocess for a group of machines.
func (s *Step) Wait(e *Exp) error {
	err := s.g.Wait(s.CmdID)
	if err == nil {
		e.mux.Lock()
		delete(e.as, s.ID)
		e.mux.Unlock()
	}
	return err
}

// Stop will signal with SIGKILL or SIGINT (default) a subprocess for a group of machines and
// then wait on the results.
func (s *Step) Stop(e *Exp) error {
	var err error
	if s.Kill {
		err = s.g.Kill(s.CmdID)
	} else {
		err = s.g.Stop(s.CmdID)
	}
	if err == nil {
		e.mux.Lock()
		delete(e.as, s.ID)
		e.mux.Unlock()
	}
	return err
}

// Save saves the results of a subprocess for a group.
func (s *Step) Save() error {
	return s.g.Save(s.CmdID, s.ID)
}

// crash is used to try to gracefully close out a step when something goes
// wrong.
func (s *Step) crash() {
	if s.g == nil {
		return
	}
	s.g.Kill(s.CmdID)
}
